from coin.blockchain.cosmos.native_private_client import \
    CosmosNativePrivateClient as nprivc
from coin.blockchain.cosmos.constants import (
    COSMOS_TOKEN_TO_ADDRESS, OSMOSIS_TOKEN_TO_ADDRESS, NODE_SERVER, TOKEN_TO_UNIT)
from coin.exchange.base.kr_rest.private_client_base import \
    ExWalletPrivateClientBase
from coin.blockchain.request_util import (
    retry_error, parse_currency_balance)


class CosmosPrivateClient(ExWalletPrivateClientBase):
  def __init__(self, *, blockchain_name, node_info=None, api_key=None):
    self._blockchain_name = blockchain_name
    self._token_unit = TOKEN_TO_UNIT
    node_info = node_info or NODE_SERVER[self._blockchain_name]
    self._nprivc = nprivc(node_info=node_info, api_key=api_key)
    if self._blockchain_name == 'cosmos':
      self._token_to_address = COSMOS_TOKEN_TO_ADDRESS
    elif self._blockchain_name == 'osmosis':
      self._token_to_address = OSMOSIS_TOKEN_TO_ADDRESS
    else:
      raise ValueError(self._blockchain_name)

  def get_token_balance(self, *, address, token_address, end_height):
    return self._nprivc.get_token_balance(address, token_address, end_height)

  def query_account_balance_impl(self, *, address, currency, end_height=None):
    token_address = self._token_to_address[currency]
    balance = retry_error(
      self.get_token_balance,
      address=address,
      token_address=token_address,
      end_height=end_height)
    balance = parse_currency_balance(currency, self._token_unit[currency], balance)
    return balance

  def get_latest_block_number_timestamp(self):
    return retry_error(self._nprivc.get_latest_block_number_timestamp)

  def get_block_timestamp(self, block_number):
    return retry_error(self._nprivc.get_block_timestamp, block_number=block_number)


if __name__ == '__main__':
  wallet_address = 'osmo16qwnqs6ywuf3kzd75cey4d2htmtkw0jnjvuwny'
  osmo_client = CosmosPrivateClient(blockchain_name='osmo')
  block_num, block_ts = osmo_client.get_latest_block_number_timestamp()
  print(osmo_client.query_account_balance_impl(address=wallet_address, currency='OSMO'))
  print(osmo_client.query_account_balance_impl(address=wallet_address, currency='USDC'))
